package com.github.shiqiyue.app.common.trace;

import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;

/***
 * 异步任务装饰器,支持链路信息传递
 * @author wenyao.wu
 */
public class TraceTaskDecorator implements TaskDecorator {

    @Override
    public Runnable decorate(Runnable runnable) {
        String traceId = MDC.get(TraceConstants.TRACE_ID);
        String memberId = MDC.get(TraceConstants.MEMBER_ID);
        return () -> {
            try {
                MDC.put(TraceConstants.TRACE_ID, traceId);
                MDC.put(TraceConstants.MEMBER_ID, memberId);
                runnable.run();
            } catch (Exception e) {
                throw e;
            } finally {
                MDC.remove(TraceConstants.TRACE_ID);
                MDC.remove(TraceConstants.MEMBER_ID);
            }
        };
    }
}
